package com.example.scribe.backend;

import com.example.scribe.SyntaxNode;
import java.util.List;

/**
 * Utility class for AST traversal and node finding operations.
 * Provides common methods used across multiple handlers.
 */
public final class ASTUtils {
    private ASTUtils() {
        // Utility class, prevent instantiation
    }

    /**
     * Find the first node of the specified type in the subtree.
     * 
     * @param node the root node to search from
     * @param type the type to search for
     * @return the first matching node, or null if not found
     */
    public static SyntaxNode findFirst(SyntaxNode node, String type) {
        if (node == null) return null;
        if (type.equals(node.type())) return node;
        for (SyntaxNode c : node.children()) {
            SyntaxNode res = findFirst(c, type);
            if (res != null) return res;
        }
        return null;
    }

    /**
     * Find the nth occurrence of a node type in the subtree.
     * 
     * @param node the root node to search from
     * @param type the type to search for
     * @param n the zero-based index of the occurrence to find
     * @return the nth matching node, or null if not found
     */
    public static SyntaxNode findNth(SyntaxNode node, String type, int n) {
        int[] count = new int[] { 0 };
        return findNthRec(node, type, n, count);
    }

    private static SyntaxNode findNthRec(SyntaxNode node, String type, int n, int[] count) {
        if (node == null) return null;
        if (type.equals(node.type())) {
            if (count[0] == n) return node;
            count[0]++;
        }
        for (SyntaxNode c : node.children()) {
            SyntaxNode res = findNthRec(c, type, n, count);
            if (res != null) return res;
        }
        return null;
    }

    /**
     * Collect all nodes in the subtree into a list.
     * 
     * @param node the root node
     * @param result the list to collect nodes into
     */
    public static void collectAllNodes(SyntaxNode node, List<SyntaxNode> result) {
        if (node == null) return;
        result.add(node);
        for (SyntaxNode child : node.children()) {
            collectAllNodes(child, result);
        }
    }

    /**
     * Check if the subtree contains any node of the specified type.
     * 
     * @param node the root node to search from
     * @param type the type to search for
     * @return true if a node of the specified type exists
     */
    public static boolean containsType(SyntaxNode node, String type) {
        return findFirst(node, type) != null;
    }

    /**
     * Extract identifier name from a node (handles various node structures).
     * 
     * @param node the node to extract from
     * @return the identifier name, or null if not found
     */
    public static String extractIdentifierName(SyntaxNode node) {
        if (node == null) return null;
        String nodeType = node.type();
        if ("Name".equals(nodeType) || "Identifier".equals(nodeType) || "identifier".equals(nodeType)) {
            return node.text();
        }
        // Search in children
        for (SyntaxNode c : node.children()) {
            String n = extractIdentifierName(c);
            if (n != null) return n;
        }
        return null;
    }

    /**
     * Extract expression node from a return statement.
     * 
     * @param blockStmt the block statement to search in
     * @return the expression node, or null if not found
     */
    public static SyntaxNode extractReturnExpr(SyntaxNode blockStmt) {
        if (blockStmt == null) return null;
        for (SyntaxNode s : blockStmt.children()) {
            if ("statement".equals(s.type())) {
                for (SyntaxNode ch : s.children()) {
                    if ("return_stmt".equals(ch.type())) {
                        for (SyntaxNode e : ch.children()) {
                            if ("expression".equals(e.type())) return e;
                        }
                    }
                }
            }
        }
        return null;
    }

    /**
     * Find the first Call node in the subtree.
     * 
     * @param node the root node to search from
     * @return the first Call node, or null if not found
     */
    public static SyntaxNode findFirstCall(SyntaxNode node) {
        if (node == null) return null;
        if ("Call".equals(node.type())) return node;
        for (SyntaxNode c : node.children()) {
            SyntaxNode found = findFirstCall(c);
            if (found != null) return found;
        }
        return null;
    }

    /**
     * Check if a node represents an identifier (case-insensitive).
     * 
     * @param node the node to check
     * @return true if the node is an identifier
     */
    public static boolean isIdentifier(SyntaxNode node) {
        if (node == null) return false;
        String type = node.type();
        return "identifier".equals(type) || "Identifier".equals(type);
    }

    /**
     * Check if a node represents a number literal.
     * 
     * @param node the node to check
     * @return true if the node is a number
     */
    public static boolean isNumber(SyntaxNode node) {
        return node != null && "Terminal:NUMBER".equals(node.type());
    }
}

